import numpy as np
import heapq
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib as mpl

# 常量定义
M, N = 20, 30
T = 90
B_sensor = 10  # Mbps
B_receive = 100  # Mbps
total_flows = 1800
flow_data = 10  # Mb/flow
car_coverage = {0: (5, 6), 1: (15, 16), 2: (25, 26)}  # 移动车覆盖的y行
flow_start_times = [0, 30, 60]  # 数据流生成时间


def bandwidth_function(t, B, phi):
    """计算传感器上传带宽: 周期三角波函数"""
    pt = (t + phi) % 10
    if pt < 5:
        return B * (pt / 5)
    else:
        return B * ((10 - pt) / 5)


def car_x_position(t):
    """计算移动车覆盖的x起始位置: 匀速移动"""
    # 单程100秒，x范围1~20 → 速度 = (20-1)/100 = 0.19
    return 1 + int(0.19 * t)  # 覆盖[x_start, x_start+2]


class Flow:
    """数据流类"""

    def __init__(self, fid, start_time, source_node):
        self.fid = fid
        self.start_time = start_time
        self.source_node = source_node
        self.remaining = flow_data
        self.current_node = source_node
        self.target_car = self.assign_target_car()
        self.path = [source_node]  # 初始化路径包含起点
        self.finish_time = None
        self.transmissions = []  # 记录每秒传输情况

    def assign_target_car(self):
        """分配目标移动车: 基于y坐标最近原则"""
        y = self.source_node[1]
        distances = [abs(y - (coverage[0] + coverage[1]) / 2) for coverage in car_coverage.values()]
        return np.argmin(distances)


class NetworkScheduler:
    """网络调度器"""

    def __init__(self):
        self.flows = []
        self.active_flows = []
        self.initialize_flows()

        # 初始化覆盖节点的带宽参数
        self.node_B = {}
        self.node_phi = {}
        for car_id, (y1, y2) in car_coverage.items():
            for y in [y1, y2]:
                for x in range(1, M + 1):
                    self.node_B[(x, y)] = np.random.uniform(5, 15)  # 峰值带宽
                    self.node_phi[(x, y)] = np.random.uniform(0, 10)  # 相位

        # 链路结构: 仅存储右和下方向
        self.edges = []
        for x in range(1, M + 1):
            for y in range(1, N + 1):
                if x < M:
                    self.edges.append(((x, y), (x + 1, y)))
                if y < N:
                    self.edges.append(((x, y), (x, y + 1)))

    def initialize_flows(self):
        """初始化数据流: 每个传感器在指定时间生成流"""
        fid = 0
        for start_time in flow_start_times:
            for x in range(1, M + 1):
                for y in range(1, N + 1):
                    self.flows.append(Flow(fid, start_time, (x, y)))
                    fid += 1

    def dijkstra_path(self, source, target_nodes, mesh_bw, t):
        """改进的Dijkstra算法: 考虑实时链路利用率"""
        # 如果源节点已经在目标节点中，直接返回路径
        if source in target_nodes:
            return [source]

        dist = {node: float('inf') for node in [(x, y) for x in range(1, M + 1) for y in range(1, N + 1)]}
        prev = {}
        dist[source] = 0
        heap = [(0, source)]

        while heap:
            d, node = heapq.heappop(heap)
            if d > dist[node]:
                continue
            if node in target_nodes:
                path = []
                while node in prev:
                    path.append(node)
                    node = prev[node]
                path.append(source)
                return path[::-1]

            # 获取邻居节点
            neighbors = []
            x, y = node
            if x > 1: neighbors.append((x - 1, y))
            if x < M: neighbors.append((x + 1, y))
            if y > 1: neighbors.append((x, y - 1))
            if y < N: neighbors.append((x, y + 1))

            for nb in neighbors:
                # 计算链路权重: 基于剩余带宽的倒数 (带宽越紧张权重越高)
                edge = tuple(sorted((node, nb)))
                used_bw = max(0.1, B_sensor - mesh_bw[edge])  # 避免除以零
                weight = 1 / used_bw
                new_dist = d + weight
                if new_dist < dist[nb]:
                    dist[nb] = new_dist
                    prev[nb] = node
                    heapq.heappush(heap, (new_dist, nb))
        return None  # 无路径

    def schedule(self):
        """主调度循环: 每秒更新并分配资源"""
        for t in range(T):
            # 1. 更新移动车位置
            x_start = car_x_position(t)

            # 2. 加入新流 (避免重复添加)
            if t in flow_start_times:
                for flow in self.flows:
                    if flow.start_time == t and flow not in self.active_flows:
                        self.active_flows.append(flow)

            # 3. 初始化带宽资源
            mesh_bw = {edge: B_sensor for edge in self.edges}
            car_bw = [B_receive] * 3  # 每辆车的剩余接收带宽

            # 计算覆盖节点的上传带宽
            upload_bw = {}
            for car_id, (y1, y2) in car_coverage.items():
                for y in [y1, y2]:
                    for x in range(x_start, x_start + 3):
                        if 1 <= x <= M:
                            node = (x, y)
                            upload_bw[node] = bandwidth_function(t, self.node_B[node], self.node_phi[node])

            # 4. 处理活动流
            temp_completed = []  # 临时存储本时步完成的流
            for flow in self.active_flows:
                if flow.remaining <= 0:
                    temp_completed.append(flow)
                    continue

                # 确定目标移动车覆盖的节点
                target_y = car_coverage[flow.target_car]
                target_nodes = [
                    (x, y) for y in target_y
                    for x in range(x_start, x_start + 3)
                    if 1 <= x <= M
                ]

                # 检查当前节点是否在覆盖范围内
                if flow.current_node in target_nodes:
                    # 直接上传
                    bw_available = min(
                        upload_bw.get(flow.current_node, 0),
                        car_bw[flow.target_car],
                        flow.remaining
                    )
                    if bw_available > 0:
                        # 记录传输
                        flow.remaining -= bw_available
                        # 更新带宽资源
                        upload_bw[flow.current_node] -= bw_available
                        car_bw[flow.target_car] -= bw_available
                        flow.transmissions.append((
                            t, flow.current_node, flow.target_car, bw_available
                        ))
                        if flow.remaining <= 0:
                            flow.finish_time = t
                            temp_completed.append(flow)
                    continue

                # 动态路径规划
                path = self.dijkstra_path(flow.current_node, target_nodes, mesh_bw, t)
                if not path or len(path) < 2:
                    continue  # 无有效路径

                next_node = path[1]  # 下一跳节点
                edge = tuple(sorted((flow.current_node, next_node)))

                # 计算可用带宽
                bw_available = min(mesh_bw.get(edge, 0), flow.remaining)
                if bw_available > 0:
                    # 更新状态
                    flow.remaining -= bw_available
                    mesh_bw[edge] -= bw_available
                    flow.current_node = next_node  # 移动到下一节点
                    flow.path.append(next_node)  # 记录路径
                    flow.transmissions.append((
                        t, next_node, None, bw_available
                    ))
                    # 设置完成时间
                    if flow.remaining <= 0:
                        flow.finish_time = t
                        temp_completed.append(flow)

            # 移除本时步完成的流
            for flow in temp_completed:
                if flow in self.active_flows:
                    self.active_flows.remove(flow)

    def calculate_metrics(self):
        """计算评价指标"""
        total_data = total_flows * flow_data
        received_data = total_data - sum(f.remaining for f in self.flows)
        loss_rate = 1 - received_data / total_data

        # 只计算已完成流量的时延
        completed_flows = [f for f in self.flows if f.finish_time is not None]
        if not completed_flows:
            avg_delay = 0
        else:
            delays = [f.finish_time - f.start_time + 1 for f in completed_flows]
            avg_delay = np.mean(delays)

        # 打印统计信息
        print(f"已完成流量: {len(completed_flows)}/{total_flows}")
        print(f"总丢包量: {total_data - received_data:.2f} Mb")
        return loss_rate, avg_delay

    def visualize_flow_path(self, fid):
        """可视化指定流的传输路径"""
        flow = self.flows[fid]
        if not flow.transmissions:
            print(f"Flow {fid} not transmitted")
            return

        fig, ax = plt.subplots(figsize=(12, 8))

        # 绘制网格
        for x in range(1, M + 1):
            for y in range(1, N + 1):
                ax.plot(x, y, 'ko', markersize=4)

        # 标记移动车覆盖区域 - 修复颜色格式
        colors = ['red', 'green', 'blue']
        for t in range(T):
            x_start = car_x_position(t)
            for car_id, (y1, y2) in car_coverage.items():
                for y in [y1, y2]:
                    for x in range(x_start, x_start + 3):
                        if 1 <= x <= M:
                            # 使用 RGBA 格式表示颜色（最后一个值0.1表示透明度）
                            rgba = mpl.colors.to_rgba(colors[car_id], alpha=0.1)
                            ax.add_patch(patches.Rectangle(
                                (x - 0.5, y - 0.5), 1, 1,
                                linewidth=0.5,
                                edgecolor=colors[car_id],
                                facecolor=rgba  # 使用正确的 RGBA 格式
                            ))

        # 绘制流路径
        path_x = [node[0] for node in flow.path]
        path_y = [node[1] for node in flow.path]
        ax.plot(path_x, path_y, 'm-', linewidth=2, marker='o', markersize=6)

        # 标记起点和终点
        ax.plot(flow.source_node[0], flow.source_node[1], 'go', markersize=8, label='Start')
        if flow.finish_time is not None:
            last_node = flow.path[-1]
            ax.plot(last_node[0], last_node[1], 'ro', markersize=8, label='End')

        ax.set_title(f"Flow Path (ID: {fid})")
        ax.set_xlabel("X Coordinate")
        ax.set_ylabel("Y Coordinate")
        ax.legend()
        plt.grid(True)
        plt.show()


# 执行调度
scheduler = NetworkScheduler()
scheduler.schedule()
loss_rate, avg_delay = scheduler.calculate_metrics()

print(f"丢包率: {loss_rate:.4f}")
print(f"平均时延: {avg_delay:.2f} seconds")

# 可视化示例流 (ID=0)
scheduler.visualize_flow_path(0)